#使用已训练的神经网络实现MNIST数据集的推理处理
import sys, os
sys.path.append(os.pardir)
import numpy as np
from dataset.mnist import load_mnist
import pickle

def get_data():  #获取MNIST数据集的测试数据集
    (x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, flatten=True, one_hot_label=False)
    return x_test, t_test

def init_network():  #初始化神经网络
    with open('sample_weight.pkl', 'rb') as f:
        network = pickle.load(f)
    return network

def sigmoid(x):  #定义sigmoid函数
    return 1 / (1 + np.exp(-x))

def softmax(a):  #定义softmax函数，添加溢出对策
    c = np.max(a)
    exp_a = np.exp(a - c)
    sum_exp_a = np.sum(exp_a)
    y = exp_a / sum_exp_a
    return y

def predict(network, x):  #推理函数
    W1, W2, W3 = network['W1'], network['W2'], network['W3']
    b1, b2, b3 = network['b1'], network['b2'], network['b3']
    a1 = np.dot(x, W1) + b1
    z1 = sigmoid(a1)
    a2 = np.dot(z1, W2) + b2
    z2 = sigmoid(a2)
    a3 = np.dot(z2, W3) + b3
    y = softmax(a3)
    return y

#计算识别精度
x, t = get_data()
network = init_network()
accuracy_cnt = 0
for i in range(len(x)):
    y = predict(network, x[i])
    p = np.argmax(y)  #获取概率最高的元素的索引
    if p == t[i]:
        accuracy_cnt += 1

print('识别精度为： ' + str(float(accuracy_cnt) / len(x)))

#基于批处理的代码实现(计算识别精度）
x, t = get_data()
network = init_network()
batch_size = 100  #定义批数量
accuracy_cnt = 0
for i in range(0, len(x), batch_size):
    x_batch = x[i:i+batch_size]
    y_batch = predict(network, x_batch)
    p = np.argmax(y_batch, axis=1)
    accuracy_cnt += np.sum(p == t[i:i+batch_size])
print('识别精度为： ' + str(float(accuracy_cnt) / len(x)))

